# -*- coding: utf-8 -*-
import os
import xml.etree.ElementTree as ET
import subprocess
import time
import matplotlib.pyplot as plt


class JMeterAutoTest:
    def __init__(self, jmeter_path, jmx_template, output_dir):
        self.jmeter_path = jmeter_path
        self.jmx_template = jmx_template
        self.output_dir = output_dir
        os.makedirs(output_dir, exist_ok=True)

    def modify_thread_count(self, new_count, output_jmx):
        """修改JMX文件中的线程数"""
        tree = ET.parse(self.jmx_template)
        root = tree.getroot()

        # 找到线程组元素
        for thread_group in root.iter('ThreadGroup'):
            # 修改线程数
            for elem in thread_group:
                if elem.tag == 'stringProp' and elem.attrib.get('name') == 'ThreadGroup.num_threads':
                    elem.text = str(new_count)

        # 保存修改后的文件
        tree.write(output_jmx, encoding='utf-8', xml_declaration=True)
        print(f"已修改线程数为: {new_count}，保存到: {output_jmx}")

    def run_test(self, jmx_file, test_name):
        """执行JMeter测试"""
        timestamp = time.strftime("%Y%m%d_%H%M%S")
        jtl_file = os.path.join(self.output_dir, f"{test_name}_{timestamp}.jtl")
        report_dir = os.path.join(self.output_dir, f"report_{test_name}_{timestamp}")

        cmd = [
            self.jmeter_path,
            '-n',  # 非GUI模式
            '-t', jmx_file,
            '-l', jtl_file,
            '-e',  # 生成HTML报告
            '-o', report_dir
        ]

        print(f"执行命令: {' '.join(cmd)}")
        process = subprocess.Popen(cmd, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
        stdout, stderr = process.communicate()

        if process.returncode == 0:
            print("测试执行成功")
            return jtl_file
        else:
            print(f"测试执行失败: {stderr.decode('utf-8')}")
            return None

    def analyze_results(self, jtl_file):
        """简单分析结果(可扩展)"""
        print(f"测试结果文件: {jtl_file}")
        return True

    def run_with_different_threads(self, thread_counts, test_name_prefix):
        """使用不同线程数运行测试"""
        results = []

        for count in thread_counts:
            print(f"\n开始测试线程数: {count}")
            modified_jmx = os.path.join(self.output_dir, f"modified_{count}.jmx")

            # 1. 修改线程数
            self.modify_thread_count(count, modified_jmx)

            # 2. 执行测试
            test_name = f"{test_name_prefix}_{count}threads"
            jtl_file = self.run_test(modified_jmx, test_name)

            if jtl_file:
                # 3. 分析结果
                success = self.analyze_results(jtl_file)
                results.append((count, success, jtl_file))

        return results

    def visualize_results(self, results):
        """可视化测试结果(示例)"""
        counts = [r[0] for r in results]
        plt.plot(counts, 'o-')
        plt.title('Thread Count Variation')
        plt.xlabel('Test Sequence')
        plt.ylabel('Thread Count')
        plt.show()


if __name__ == "__main__":
    # 配置参数
    JMETER_PATH = r"D:\apache-jmeter-5.6.3\bin\jmeter.bat"  # 替换为你的JMeter可执行文件路径
    TEMPLATE_JMX = r"D:\jmeter\轻商城项目性能测试脚本.jmx"  # 使用完整路径
    OUTPUT_DIR = "test_results"

    # 初始化测试器
    tester = JMeterAutoTest(JMETER_PATH, TEMPLATE_JMX, OUTPUT_DIR)

    # 定义要测试的线程数列表
    THREAD_COUNTS = [50, 100, 150, 200]  # 可以调整为你的需求

    # 执行测试
    test_results = tester.run_with_different_threads(THREAD_COUNTS, "performance_test")

    # 可视化结果
    tester.visualize_results(test_results)